import os
from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, Dict, Any
from dotenv import load_dotenv
import base64

from qcloud_hunyuan3d import QCloudHunyuan3D
from db_manager import DatabaseManager
# 导入3D模型评估API路由
from three_d_effect_evaluation.api import router as evaluation_router


# 加载环境变量
load_dotenv()

# 创建FastAPI应用
app = FastAPI(
    title="3D模型生成服务",
    description="通过腾讯云混元3D API实现的文本/图片到3D模型生成服务，包含模型评估功能",
    version="1.0.0"
)

# 配置CORS中间件
# app.add_middleware(
#     CORSMiddleware,
#     allow_origins=["*"],
#     allow_credentials=True,
#     allow_methods=["*"],
#     allow_headers=["*"],
# )

# 初始化腾讯云混元3D客户端
hunyuan3d_client = QCloudHunyuan3D()

# 初始化数据库管理器
db_manager = DatabaseManager()

# 注册3D模型评估API路由
app.include_router(evaluation_router)

# 请求模型定义
class TextTo3DRequest(BaseModel):
    prompt: str = Field(..., description="3D模型的文本描述")
    user_id: int = Field(..., description="用户ID")
    options: Optional[Dict[str, Any]] = Field(None, description="可选参数")

class ImageTo3DURLRequest(BaseModel):
    image_url: str = Field(..., description="图片URL或base64编码")
    user_id: int = Field(..., description="用户ID")
    options: Optional[Dict[str, Any]] = Field(None, description="可选参数")

# 健康检查接口
@app.get("/health", tags=["健康检查"])
async def health_check():
    """服务健康检查接口"""
    return {
        "status": "ok",
        "message": "3D模型生成服务运行正常",
        "version": "1.0.0"
    }

# 根据文本生成3D模型接口
@app.post("/api/v2/generate/text-to-3d", tags=["3D模型生成"])
async def generate_3d_from_text(request: TextTo3DRequest):
    """
    根据文本描述生成3D模型
    
    - **prompt**: 3D模型的文本描述，例如 "一匹马"、"一个茶杯" 等
    - **user_id**: 用户ID，用于数据库关联
    - **options**: 可选参数，用于调整生成参数，可包含action字段指定任务类型
    
    返回包含任务ID的响应，可用于后续查询任务状态
    """
    try:
        if not request.prompt or request.prompt.strip() == "":
            raise HTTPException(status_code=400, detail="文本描述不能为空")
        
        result = hunyuan3d_client.generate_3d_from_text(
            prompt=request.prompt,
            options=request.options
        )
        
        # 检查响应格式是否符合预期
        if not result or 'Response' not in result:
            raise HTTPException(status_code=500, detail="API返回格式异常")
        
        if 'JobId' in result['Response']:
                job_id = result['Response']['JobId']
                request_id = result['Response'].get('RequestId', '')
                
                # 插入数据库记录
                db_manager.insert_generation_record(
                    user_id=request.user_id,
                    generation_method='prompt',
                    prompt=request.prompt,
                    task_id=job_id,
                    request_id=request_id,
                    status='RUN'
                )
            
                response_data = {
                "job_id": job_id,
                "request_id": request_id
            }
        else:
            # 如果API返回了其他格式，保留原始数据
            response_data = result['Response']
        
        return {
            "code": 1,
            "msg": "success",
            "data": {
                "success": True,
                "data": response_data
            }
        }
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"生成3D模型时发生错误: {str(e)}"
        )

# 根据图片URL生成3D模型接口
@app.post("/api/v2/generate/image-to-3d/url", tags=["3D模型生成"])
async def generate_3d_from_image_url(request: ImageTo3DURLRequest):
    """
    根据图片URL或base64编码生成3D模型
    
    - **image_url**: 图片的URL地址或base64编码字符串
    - **user_id**: 用户ID，用于数据库关联
    - **options**: 可选参数，用于调整生成参数，可包含action字段指定任务类型
    
    返回包含任务ID的响应，可用于后续查询任务状态
    """
    try:
        if not request.image_url or request.image_url.strip() == "":
            raise HTTPException(status_code=400, detail="图片URL不能为空")
        
        result = hunyuan3d_client.generate_3d_from_image(
            image_url=request.image_url,
            options=request.options
        )
        
        # 检查响应格式是否符合预期
        if not result or 'Response' not in result:
            raise HTTPException(status_code=500, detail="API返回格式异常")
        
        # 确保返回的是JobId而不是TaskId
        if 'JobId' in result['Response']:
            job_id = result['Response']['JobId']
            request_id = result['Response'].get('RequestId', '')
            
            # 插入数据库记录
            db_manager.insert_generation_record(
                user_id=request.user_id,
                generation_method='image',
                image_url=request.image_url,
                task_id=job_id,
                request_id=request_id,
                status='RUN'
            )
            
            response_data = {
                "job_id": job_id,
                "request_id": request_id
            }
        else:
            # 如果API返回了其他格式，保留原始数据
            response_data = result['Response']
        
        return {
            "code": 1,
            "msg": "success",
            "data": {
                "success": True,
                "data": response_data
            }
        }
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"根据图片生成3D模型时发生错误: {str(e)}"
        )

# 根据图片文件生成3D模型接口
@app.post("/api/v2/generate/image-to-3d/upload", tags=["3D模型生成"])
async def generate_3d_from_image_upload(
    user_id: int = Form(...),
    image: UploadFile = File(...),
    options: Optional[str] = Form(None)
):
    """
    上传图片文件生成3D模型
    
    - **user_id**: 用户ID，用于数据库关联
    - **image**: 要上传的图片文件
    - **options**: 可选参数（JSON字符串格式），用于调整生成参数，可包含action字段指定任务类型
    
    返回包含任务ID的响应，可用于后续查询任务状态
    """
    try:
        # 读取图片文件内容并转换为base64编码
        image_content = await image.read()
        image_base64 = base64.b64encode(image_content).decode("utf-8")
        
        # 解析options参数
        options_dict = None
        if options:
            try:
                import json
                options_dict = json.loads(options)
            except json.JSONDecodeError:
                raise HTTPException(status_code=400, detail="options参数格式不正确，应为JSON字符串")
        
        result = hunyuan3d_client.generate_3d_from_image(
            image_url=image_base64,
            options=options_dict
        )
        
        # 检查响应格式是否符合预期
        if not result or 'Response' not in result:
            raise HTTPException(status_code=500, detail="API返回格式异常")
        
        # 确保返回的是JobId而不是TaskId
        if 'JobId' in result['Response']:
            job_id = result['Response']['JobId']
            request_id = result['Response'].get('RequestId', '')
            
            # 插入数据库记录
            db_manager.insert_generation_record(
                user_id=user_id,
                generation_method='image',
                image_url=f"Uploaded image: {image.filename}",  # 记录上传文件名作为image_url
                task_id=job_id,
                request_id=request_id,
                status='RUN'
            )
            
            response_data = {
                "job_id": job_id,
                "request_id": request_id
            }
        else:
            # 如果API返回了其他格式，保留原始数据
            response_data = result['Response']
        
        return {
            "code": 1,
            "msg": "success",
            "data": {
                "success": True,
                "data": response_data
            }
        }
        
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"上传图片生成3D模型时发生错误: {str(e)}"
        )

# 查询任务状态接口
@app.get("/api/v2/task/status/{job_id}", tags=["任务管理"])
async def get_task_status(job_id: str, requestId: str, options: Optional[str] = None):
    """
    查询3D模型生成任务的状态
    
    - **job_id**: 任务ID，由生成接口返回
    - **options**: 可选参数（JSON字符串格式），可包含action字段指定查询接口类型
    
    返回任务的当前状态和相关信息
    """
    try:
        if not job_id or job_id.strip() == "":
            raise HTTPException(status_code=400, detail="任务ID不能为空")
        
        # 解析options参数
        options_dict = None
        if options:
            try:
                import json
                options_dict = json.loads(options)
            except json.JSONDecodeError:
                raise HTTPException(status_code=400, detail="options参数格式不正确，应为JSON字符串")
        
        result = hunyuan3d_client.get_task_status(task_id=job_id, options=options_dict)
        
        # 检查响应格式是否符合预期
        if not result or 'Response' not in result:
            raise HTTPException(status_code=500, detail="API返回格式异常")
        
        # 标准化状态信息，确保符合腾讯云官方文档定义
        # 任务状态：WAIT（等待中）、RUN（执行中）、FAIL（任务失败）、DONE（任务成功）
        response_data = result['Response']

        # 替换response_data中的requestId为传入的requestId
        response_data['RequestId'] = requestId
        
        # 增强任务成功时的返回信息
        if 'Status' in response_data and response_data['Status'] == 'DONE' and 'ResultFile3Ds' in response_data:
            # 确保结果文件信息格式一致
            for file_info in response_data['ResultFile3Ds']:
                # 确保必要字段存在
                if 'Type' not in file_info:
                    file_info['Type'] = 'unknown'
                if 'Url' not in file_info:
                    file_info['Url'] = ''
        
        # 增强任务失败时的错误信息
        if 'Status' in response_data and response_data['Status'] == 'FAIL':
            if 'ErrorCode' not in response_data:
                response_data['ErrorCode'] = 'UnknownError'
            if 'ErrorMessage' not in response_data:
                response_data['ErrorMessage'] = '任务执行失败，但未返回具体错误信息'
        
        return {
            "code": 1,
            "msg": "success",
            "data": {
                "success": True,
                "data": response_data
            }
        }
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"查询任务状态时发生错误: {str(e)}"
        )

# 启动应用程序的入口点
if __name__ == "__main__":
    import uvicorn
    
    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", 8080))
    
    print(f"3D模型生成服务启动中...")
    print(f"服务将运行在: http://{host}:{port}")
    print(f"API文档: http://{host}:{port}/docs")
    print(f"注意：请确保在.env文件中配置了正确的腾讯云API密钥")
    
    uvicorn.run("main:app", host=host, port=port, reload=True)